import orekit
from orekit.pyhelpers import setup_orekit_curdir
from org.orekit.models.earth.atmosphere import HarrisPriester
from org.orekit.propagation.conversion import DormandPrince853IntegratorBuilder
from org.orekit.propagation.numerical import NumericalPropagator
from org.orekit.propagation.analytical.tle import TLE, TLEPropagator
from org.orekit.bodies import CelestialBodyFactory
from org.orekit.forces.gravity import HolmesFeatherstoneAttractionModel
from org.orekit.forces.gravity.potential import GravityFieldFactory
from org.orekit.forces.drag import DragForce, IsotropicDrag
from org.orekit.forces.radiation import SolarRadiationPressure
from org.orekit.orbits import KeplerianOrbit, OrbitType
from org.orekit.frames import FramesFactory
from org.orekit.utils import Constants, IERSConventions
from org.orekit.time import AbsoluteDate, TimeScalesFactory
import math

# 初始化 Orekit 环境
orekit.initVM()
setup_orekit_curdir("../tools/orekit-data-master")  # 替换为你的 orekit 数据路径

if __name__ == "__main__":
    # 1. 解析 TLE 数据
    tle_line1 = "1 27944U 03042F   23365.91886760  .00000101  00000-0  29022-4 0  9990"
    tle_line2 = "2 27944  98.2882 140.0552 0011603 327.1746  32.8733 14.63434870 81963"
    tle = TLE(tle_line1, tle_line2)

    # 2. 获取初始轨道状态
    tle_propagator = TLEPropagator.selectExtrapolator(tle)
    initial_state = tle_propagator.propagate(tle.getDate())

    # 3. 创建数值传播器
    # Orekit 12.2 的正确构建方式
    min_step = 0.1
    max_step = 300.0
    position_tolerance = 10.0

    # 关键修正：Orekit 12.2 使用新的构建方式
    integrator = DormandPrince853IntegratorBuilder(
        min_step, max_step, position_tolerance
    ).buildIntegrator(initial_state.getOrbit(), PositionAngleType.MEAN)

    propagator = NumericalPropagator(integrator)
    propagator.setInitialState(initial_state)

    # 4. 添加摄动力模型（与之前相同）
    gravity_provider = GravityFieldFactory.getNormalizedProvider(50, 50)
    gravity_model = HolmesFeatherstoneAttractionModel(
        FramesFactory.getITRF(IERSConventions.IERS_2010, False),
        gravity_provider
    )
    propagator.addForceModel(gravity_model)

    atmosphere = HarrisPriester()
    drag_force = DragForce(atmosphere, IsotropicDrag(10.0, 2.2))
    propagator.addForceModel(drag_force)

    solar_radiation_pressure = SolarRadiationPressure(
        CelestialBodyFactory.getSun(),
        Constants.SUN_RADIUS,
        1.2,
        10.0
    )
    propagator.addForceModel(solar_radiation_pressure)

    # 5. 传播轨道
    target_time = tle.getDate().shiftedBy(3600.0)  # 1小时后
    future_state = propagator.propagate(target_time)

    # 6. 输出结果
    future_orbit = OrbitType.KEPLERIAN.convertType(future_state.getOrbit())

    print("\n=== 轨道预测结果 ===")
    print(f"预测时间: {target_time}")

    pos = future_state.getPVCoordinates().getPosition()
    print(f"\n位置 (m): X={pos.getX():.2f}, Y={pos.getY():.2f}, Z={pos.getZ():.2f}")

    print("\n开普勒轨道根数:")
    print(f"半长轴 (km): {future_orbit.getA() / 1000:.3f}")
    print(f"离心率: {future_orbit.getE():.6f}")
    print(f"轨道倾角 (deg): {math.degrees(future_orbit.getI()):.3f}")